# coding: UTF-8
import os
import cv2
import numpy as np
import pandas as pd
import glob
from pathlib import Path
import random
import json

# def gen_n_defect_imgs(ok_img_path)
def find_template_position(template, image, method=cv2.TM_CCOEFF_NORMED):
    """
    在模板图像中查找原始图像的位置
    """
    result = cv2.matchTemplate(template, image, method)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)

    if method in [cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED]:
        top_left = min_loc
        match_val = 1 - min_val if method == cv2.TM_SQDIFF_NORMED else min_val
    else:
        top_left = max_loc
        match_val = max_val

    h, w = image.shape[:2]
    bottom_right = (top_left[0] + w, top_left[1] + h)

    return top_left, bottom_right, match_val

def gen_one(defect_patch, defect_mask, defect_cls, cropped_part_mask, ok_img, ok_img_name, index, output_dir):
    h, w = defect_patch.shape[:2]
    print(f"    patch 尺寸: {defect_patch.shape}, mask 尺寸: {defect_mask.shape}")

    mold_mask_bool = (cropped_part_mask == 255).astype(np.uint8)

    kernel = np.ones((h, w), dtype=np.uint8)
    valid_mask = cv2.erode(mold_mask_bool, kernel)

    # valid_centers = np.argwhere(valid_mask == 1)

    # 获取所有合法的 patch 中心坐标
    all_centers = np.argwhere(valid_mask == 1)

    # 限制粘贴中心点必须落在图像中心的 800x800 区域
    h_img, w_img = ok_img.shape[:2]
    margin_x = max((w_img - 1000) // 2, 0)
    margin_y = max((h_img - 1000) // 2, 0)
    x_min = margin_x
    x_max = w_img - margin_x
    y_min = margin_y
    y_max = h_img - margin_y

    valid_centers = np.array([
        [y, x] for y, x in all_centers
        if x_min <= x <= x_max and y_min <= y <= y_max
    ])

    if len(valid_centers) == 0:
        print("    ❌ 没有足够空间贴 patch。")
        return 

    center_y, center_x = valid_centers[random.randint(0, len(valid_centers) - 1)]

    center = (
        min(max(center_x, w // 2), ok_img.shape[1] - w // 2 - 1),
        min(max(center_y, h // 2), ok_img.shape[0] - h // 2 - 1)
    )
    print(f"    🎯 粘贴位置中心: {center}")



    try:
        ys, xs = np.where(defect_mask == 255)
        offset_x = center[0] - w // 2
        offset_y = center[1] - h // 2

        mixed = ok_img.copy()
        for y, x in zip(ys, xs):
            dst_y = y + offset_y
            dst_x = x + offset_x
            if 0 <= dst_y < mixed.shape[0] and 0 <= dst_x < mixed.shape[1]:
                mixed[dst_y, dst_x] = defect_patch[y, x]

        # 检查差异区域
        diff = cv2.absdiff(mixed, ok_img)
        diff_gray = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
        changed_area = np.count_nonzero(diff_gray)
        mask_area = np.count_nonzero(defect_mask)

        print(f"    ⏱️ 掩膜区域像素数: {mask_area}, 差异像素数: {changed_area}")

        if changed_area > 0.5 * mask_area:
            print("    ✅ 缺陷区域与原图有明显差异，确认缺陷图已贴上。")
        else:
            print("    ⚠️ 图像差异太小，疑似粘贴失败。")

    except Exception as e:
        print(f"    粘贴失败: {ok_img_name}, 错误: {e}")
        return

    save_dir = os.path.join(output_dir, defect_cls)
    os.makedirs(save_dir, exist_ok=True)

    out_img_name = Path(ok_img_name).stem + "_" + str(index) + '.jpg'
    out_json_name = Path(ok_img_name).stem + "_" + str(index) + '.json'

    cv2.imwrite(os.path.join(save_dir, out_img_name), mixed)

    # 构造 JSON 注释
    contours, _ = cv2.findContours(defect_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    shapes = []
    for cnt in contours:
        if len(cnt) < 3:
            continue
        points = cnt.squeeze().tolist()
        if not isinstance(points[0], list):
            points = [points]
        shape = {
            "label": defect_cls,
            "points": [[float(pt[0] + center[0] - w // 2), float(pt[1] + center[1] - h // 2)] for pt in points],
            "group_id": None,
            "shape_type": "polygon",
            "flags": {}
        }
        shapes.append(shape)

    json_dict = {
        "version": "5.2.0.post4",
        "flags": {},
        "shapes": shapes,
        "imagePath": out_img_name,
        "imageData": None,
        "imageHeight": ok_img.shape[0],
        "imageWidth": ok_img.shape[1]
    }

    with open(os.path.join(save_dir, out_json_name), "w") as f:
        json.dump(json_dict, f, indent=2)


def gen_defect_imgs(cut_img_path, cut_mask_path, cur_part_name, cur_code, df_ok, num, template_img, part_memo, output_root):
    defect_patch = cv2.imread(cut_img_path)
    defect_mask = cv2.imread(cut_mask_path, cv2.IMREAD_GRAYSCALE)

    df_ok = df_ok.sample(n=num)
    # df_ok = df_ok.iloc[:num]

    for index, row in df_ok.iterrows():
        ok_img_path = row['img_path']
        # code = row['code']

        ok_img = cv2.imread(ok_img_path, 1)
        ok_img_gray = cv2.cvtColor(ok_img, cv2.COLOR_BGR2GRAY)
        ok_h, ok_w = ok_img.shape[:2]

        top_left, bottom_right, match_val = find_template_position(template_img, ok_img_gray)
        x_off, y_off = top_left
        cropped_part_mask = part_memo[cur_part_name][y_off:y_off+ok_h, x_off:x_off+ok_w]

        gen_one(defect_patch, defect_mask, cur_code, cropped_part_mask, ok_img, Path(cut_img_path).name, index, output_root)


def main():

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True, nb_workers=8) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    cut_defect_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10_SD/20250530-SD-small-defect-patch'
    ok_dir = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10_SD/10X/5502W/OK'
    template_path = r"/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10_SD/template/5502W_SD/RGB.png"
    part_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10_SD/template/5502W_SD'
    output_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10_SD/20250530-SD-small-defect_generated'
    num = 10

    category_mask_suffix_map = {
            "TG05": "acom1",
            "TS28": "data",
            "TS15": "data"
    }

    #  模板相关
    scale_ratio = 1.07766
    resize_ratio = 1.0 / scale_ratio
    part_memo = {}

    # 加载并缩放模板图像
    template_img = cv2.imread(template_path, cv2.IMREAD_GRAYSCALE)
    template_img = cv2.resize(template_img, (0, 0), fx=resize_ratio, fy=resize_ratio, interpolation=cv2.INTER_AREA)
    for part_name in os.listdir(part_root):
        if part_name == "RGB.png": continue
        part_stem = Path(part_name).stem 
        part_img_resized= cv2.imread(os.path.join(part_root, part_name), cv2.IMREAD_GRAYSCALE)
        part_img_resized = cv2.resize(part_img_resized, (0, 0), fx=resize_ratio, fy=resize_ratio, interpolation=cv2.INTER_NEAREST)
        part_memo[part_stem] = part_img_resized


    df_ok = pd.DataFrame(data={"img_path":glob.glob(os.path.join(ok_dir, "*.jpg"), recursive=True)})
    # df_ok['code'] = df_ok['img_path'].apply(lambda x: Path(x).parent.name)
    # df_ok = df_ok[df_ok['code'].isin(category_mask_suffix_map.keys())]
    print(' len df_ok', len(df_ok))

    df_cut = pd.DataFrame(data={"img_path":glob.glob(os.path.join(cut_defect_root, "*/*_patch.png"), recursive=True)})
    df_cut['mask_path'] = df_cut['img_path'].apply(lambda x: os.path.join(str(Path(x).parent), str(Path(x).stem[:-5])+'mask.png'))
    df_cut['code'] = df_cut['img_path'].apply(lambda x: Path(x).parent.name)
    df_cut = df_cut[df_cut['code'].isin(category_mask_suffix_map.keys())]
    print(' len df_cut', len(df_cut))


    df_cut.parallel_apply(lambda x: gen_defect_imgs(x['img_path'], x['mask_path'], category_mask_suffix_map[x['code']], x['code'], df_ok.copy(), num, template_img, part_memo, output_root), axis=1)
    # parallel_apply


if __name__=="__main__":
    main()
